######## Functions for making container images (cpio, squashfs) ##########
from forest.registry import register_tool #, get_tool
from subprocess import Popen, PIPE #, STDOUT
import os
#import sys
import tempfile

def find(basepath):
    """ Return results similar to the Unix find command run without options

        i.e. traverse a directory tree and return all the file paths
    """

    return [os.path.normpath(os.path.join(path, file))
            for (path, dirs, files) in os.walk(basepath)
            for file in files]

try:
    import progressbar
except ImportError:
    print "Forest requires progressbar==2.2"
    raise

class Counter(progressbar.ProgressBarWidget, object):
    'Displays the current count'

    __slots__ = ('format',)

    def __init__(self, format='%d of %d'):
        self.format = format

    def update(self, pbar):
        return self.format % (pbar.currval, pbar.maxval)


progress = None
counter = 0


@register_tool('image.mkcpio')
def cpio(rootdir, outfile, compression = "gzip", compressargs = None, exclude = None):

    ''' Make a compressed CPIO archive of the given rootdir.

        rootdir - the root to cpio
        outfile - the output filename
        compression - optional - should be "xz", "gzip", "lzma", or None.
        compressargs - optional - will be used on the compression commandline.
        exclude - optional - a function matching paths in the root. (note: paths include a ./ prefix)
    '''
    if compressargs is None:
        compressargs = ["-9"]

    global progress, progfile

    if compression not in (None, "xz", "gzip", "lzma"):
        raise ValueError, "Unknown compression type %s" % compression
    chdir = lambda: os.chdir(rootdir)

    if exclude is None:
        exclude = lambda x: True

    if compression == "xz":
        compressargs.insert(0, "--check=crc32")

    if compression is None:
        compression = "cat" # this is a little silly
        compressargs = []

    # TODO: use python os.walk
    find = Popen(['find','*'], stdin = PIPE, stdout = PIPE, shell = True, preexec_fn = chdir)
    out, err = find.communicate()

    # FIXME: this breaks with symbolic links
    #with cd(rootdir):
    #    f = filter(exclude, find('.'))

    f = filter(exclude, out.split('\n'))

    widgets = [outfile, progressbar.Percentage(), ' ', progressbar.Bar(), Counter()]
    progress = progressbar.ProgressBar(widgets = widgets, maxval = len(f), term_width = 80)

    cpio_file = tempfile.SpooledTemporaryFile()
    cpio_file.write('\n'.join(f))
    cpio_file.seek(0)

    cpio = Popen(['cpio', '-o', '-c', '-v','--quiet'], bufsize = -1, stdin = cpio_file, stdout = PIPE, stderr = PIPE, preexec_fn = chdir)
    # TODO: use native gzip
    compcpio = Popen([compression] + compressargs, bufsize = -1, stdin = cpio.stdout, stdout = open(outfile, "wb"))

    counter = 0
    while cpio.returncode is None:
        if cpio.stderr:
            cpio.stderr.readline()
        cpio.poll()
        counter += 1
        if counter <= len(f):
            progress.update(counter)

    compcpio.wait()

    if compcpio.returncode != 0:
        raise Exception("Error creating cpio %s: %s" % (outfile, compcpio.returncode))

    progress.finish()
    return compcpio.returncode


if __name__ == '__main__':
    def exclude(x):
        return True

    cpio('installroot', 'test.img', compression = "gzip", compressargs = ["-9"], exclude = exclude)



